import opuslib_next

# 解码编码
class OpusCodec:
    def __init__(self, sample_rate=16000, channels=1, application=opuslib_next.APPLICATION_VOIP):
        """
        初始化Opus编解码器
        :param sample_rate: 采样率
        :param channels: 通道数
        :param application: 应用类型
        """
        self.sample_rate = sample_rate
        self.channels = channels
        self.application = application
        # 创建编码器和解码器
        self.encoder = opuslib_next.Encoder(
            sample_rate, channels, application
        )
        self.decoder = opuslib_next.Decoder(
            sample_rate, channels
        )
        
    def encode(self, pcm_data):
        """
        将PCM数据编码为Opus格式
        :param pcm_data: 16位PCM数据，numpy数组
        :return: Opus编码的字节数据
        """
        # 确保数据是16位有符号整数
        pcm_data = pcm_data.astype(np.int16)
        # 编码
        encoded_data = self.encoder.encode(pcm_data.tobytes(), len(pcm_data))
        return encoded_data
    
    def decode(self, opus_data):
        """
        解码Opus数据为PCM格式
        :param opus_data: Opus编码的字节数据
        :return: 16位PCM数据，numpy数组
        """
        # 解码
        decoded_data = self.decoder.decode(opus_data, frame_size=960)  # 960是60ms的帧大小
        # 转换为numpy数组
        pcm_array = np.frombuffer(decoded_data, dtype=np.int16)
        return pcm_array